import torch
from tqdm import tqdm
import torchvision
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader, random_split, Dataset
from torchvision.models import vgg16, VGG16_Weights, VGG16_BN_Weights
import os
from model_wrapper import ClassifyModelWrapper
from cvx2.wrapper import SplitImageClassifyModelWrapper


class MyDataset(Dataset):
    def __init__(self, dataset, transform):
        self.dataset = dataset
        self.transform = transform

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        img, label = self.dataset[idx]
        return self.transform(img), label


# 数据增强和预处理
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop((224, 224)),
        # transforms.RandomHorizontalFlip(),
		transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.48235, 0.45882, 0.40784), std=(1.0 / 255.0, 1.0 / 255.0, 1.0 / 255.0))
    ]),
    'val': transforms.Compose([
        transforms.Resize((224, 224)),
        # transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.48235, 0.45882, 0.40784), std=(1.0 / 255.0, 1.0 / 255.0, 1.0 / 255.0))
    ]),
}

# # 加载整个数据集
# dataset = ImageFolder(root='faceimg', transform=VGG16_Weights.DEFAULT.transforms())
#
# # 获取类别名称
# class_names = dataset.classes
# print(f"类别: {class_names}")

# # 数据集划分
# dataset_size = len(dataset)
# train_size = int(0.8 * dataset_size)
# val_size = dataset_size - train_size
# train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
#
# # 设置 transform
# # train_dataset.dataset.transform = data_transforms['train']
# # val_dataset.dataset.transform = data_transforms['val']
# # train_dataset = MyDataset(train_dataset, VGG16_Weights.DEFAULT.transforms())
# # val_dataset = MyDataset(val_dataset, VGG16_Weights.DEFAULT.transforms())
#
# # DataLoader
# batch_size = 16
# train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)


# 加载预训练模型并修改分类器
# model = vgg16(pretrained=True)
model = vgg16(weights=VGG16_Weights.IMAGENET1K_V1)
for param in model.features.parameters():
    param.requires_grad = False
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 5)
# model.classifier[6] = nn.Linear(num_features, len(class_names))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

train_transform = transforms.Compose([
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 亮度、对比度、饱和度和色相
    transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10)),
    VGG16_Weights.DEFAULT.transforms()
])


wrapper = SplitImageClassifyModelWrapper(model, device=device)
wrapper.train(data='faceimg', imgsz=(224, 224), transform=VGG16_Weights.DEFAULT.transforms(),
              # train_transform=train_transform,
              batch_size=32, epochs=20, lr=0.001, T_max=3)

# # wrapper = ClassifyModelWrapper(model, classes=class_names, device=device)
# # wrapper.train(train_dataset, val_dataset, lr=0.0001)
#
# # 损失函数和优化器
# criterion = nn.CrossEntropyLoss()
# optimizer = torch.optim.Adam(model.classifier.parameters(), lr=0.0001)
#
# # 训练和验证
# def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=10):
#     for epoch in range(num_epochs):
#         print(f"Epoch {epoch+1}/{num_epochs}")
#         print('-' * 10)
#
#         # 训练
#         model.train()
#         running_loss = 0.0
#         running_corrects = 0
#         for inputs, labels in tqdm(train_loader):
#             inputs, labels = inputs.to(device), labels.to(device)
#             optimizer.zero_grad()
#             outputs = model(inputs)
#             loss = criterion(outputs, labels)
#             loss.backward()
#             optimizer.step()
#             _, preds = torch.max(outputs, 1)
#             running_loss += loss.item() * inputs.size(0)
#             running_corrects += torch.sum(preds == labels.data)
#
#         epoch_loss = running_loss / len(train_loader.dataset)
#         epoch_acc = running_corrects.double() / len(train_loader.dataset)
#         print(f"训练 Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}")
#
#         # 验证
#         model.eval()
#         val_loss = 0.0
#         val_corrects = 0
#         with torch.no_grad():
#             for inputs, labels in val_loader:
#                 inputs, labels = inputs.to(device), labels.to(device)
#                 outputs = model(inputs)
#                 loss = criterion(outputs, labels)
#                 _, preds = torch.max(outputs, 1)
#                 val_loss += loss.item() * inputs.size(0)
#                 val_corrects += torch.sum(preds == labels.data)
#
#         val_loss = val_loss / len(val_loader.dataset)
#         val_acc = val_corrects.double() / len(val_loader.dataset)
#         print(f"验证 Loss: {val_loss:.4f} Acc: {val_acc:.4f}")
#
#     return model
#
# # 开始训练
# trained_model = train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=20)

# 保存模型
#torch.save(trained_model.state_dict(), 'vgg16_5class.pth')
